import os
import shutil


anno_path = '/home/sean/data/UWdevkit/Annotations/train/'
img_path = '/home/sean/data/UWdevkit/Data/train/'

dir_name = 'val'
video_path = os.path.join(img_path, dir_name)
video_anno_path = os.path.join(anno_path, dir_name)

video_len = 1327
sub_num = 13
for i in range(sub_num):
    sub_dir = dir_name+'_'+str(i).zfill(2)
    sub_video_path = os.path.join(img_path, sub_dir)
    sub_anno_path = os.path.join(anno_path, sub_dir)
    if os.path.exists(sub_video_path) or os.path.exists(sub_anno_path):
        raise Exception('sub_dir exists')
    else:
        os.mkdir(sub_video_path)
        os.mkdir(sub_anno_path)

    files = os.listdir(video_path)
    file_first_name = files[0].split('.')[0].split('_')[0]
    files_num = [int(f.split('.')[0].split('_')[-1]) for f in files]
    sorted_files_num = sorted(files_num)

    each_num = min(video_len // sub_num, len(sorted_files_num))
    selected_files = [file_first_name+'_'+str(num).zfill(4)+'.jpg' for num in sorted_files_num[:each_num]]
    selected_annos = [file_first_name+'_'+str(num).zfill(4)+'.xml' for num in sorted_files_num[:each_num]]

    for selected_file, selected_anno in zip(selected_files, selected_annos):
        absolute_jpg_path = os.path.join(video_path, selected_file)
        absolute_anno_path = os.path.join(video_anno_path, selected_anno)
        shutil.move(absolute_jpg_path, sub_video_path)
        shutil.move(absolute_anno_path, sub_anno_path)

        # print(absolute_anno_path, sub_anno_path)
    # break
